package com.yolov8.steam;

import com.yolov8.config.AppConfig;
import com.yolov8.model.FrameData;
import ai.onnxruntime.*;
import com.yolov8.util.ImageUtils;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.opencv.opencv_core.*;
import org.bytedeco.opencv.opencv_core.Mat;
import org.opencv.core.CvType;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.stereotype.Service;

import java.nio.ByteBuffer;
import java.nio.FloatBuffer;
import java.util.*;
import java.util.Arrays;
import java.util.concurrent.BlockingQueue;

import static org.bytedeco.opencv.global.opencv_core.*;
import static org.bytedeco.opencv.global.opencv_imgproc.*;

@Slf4j
public class FrameProcessor implements Runnable {

    private final OrtEnvironment env;
    private final OrtSession session;
    private final BlockingQueue<FrameData> inputQueue;
    private  BlockingQueue<Mat> outputQueue;
    private final List<String> classNames;
    private final float confidenceThreshold = 0.05f;
    private final float nmsThreshold = 0.6f;

    private final String filePath = "src/main/java/com/yolov8/data/detection.jpg";

    private final int inputWidth = 640;
    private final int inputHeight = 640;

    private volatile boolean running = true; // 添加volatile标记，确保线程可见性




    public FrameProcessor(String modelPath, BlockingQueue<FrameData> inputQueue, List<String> classNames) throws OrtException {
        this(modelPath, inputQueue, classNames, null);
    }

    public FrameProcessor(String modelPath, BlockingQueue<FrameData> pullQueue, List<String> classNames, BlockingQueue<Mat> pushQueue) throws OrtException {
        this.inputQueue = pullQueue;
        this.outputQueue = pushQueue;
        this.classNames = classNames;
        OrtSession.SessionOptions options = new OrtSession.SessionOptions();
        options.setIntraOpNumThreads(1);
//        options.enableCpuMemArena();
//        options.enableFastMath();
        this.env = OrtEnvironment.getEnvironment();
//        System.out.println("环境1："+ this.env);
        this.session = env.createSession(modelPath, options);
//        System.out.println("环境2："+ this.env);

    }


    @Override
    public void run() {
        while (running) {
            try {
                FrameData frameData = inputQueue.take();
                Mat mat = frameData.getMat();
                int originalWidth = mat.cols();
                int originalHeight = mat.rows();

                // 预处理
                Mat preprocessed = preprocessImage(mat);

//                log.info("输出预处理图像");
//                Mat saveMat = new Mat();
//                preprocessed.convertTo(saveMat, CV_8UC3, 255.0,0);
//                ImageUtils.saveMatToFile(saveMat, "src/main/java/com/yolov8/data1/preprocessed.jpg");
//                saveMat.release();

                // 转换为模型输入
                float[] inputArray = matToCHW(preprocessed);
                OnnxTensor inputTensor = OnnxTensor.createTensor(
                        env, FloatBuffer.wrap(inputArray),
                        new long[]{1, 3, inputHeight, inputWidth}
                );

                // 模型推理
                try (OrtSession.Result result = session.run(
                        Collections.singletonMap(session.getInputNames().iterator().next(), inputTensor))) {

                    // 后处理（含NMS和逻辑解析）
                    List<Detection> nmsResults = postprocessResult(result, originalWidth, originalHeight);
                    // 调用工具类绘制检测结果
                    ImageUtils.drawDetections(mat, nmsResults, classNames);

                    // 临时保存结果
                    for (Detection detection : nmsResults) {
                        if (detection.confidence > confidenceThreshold) {
                            ImageUtils.saveMatToFile(mat, filePath);
                        }
                    }
                    // 推入推流队列（clone 避免多线程并发修改）
                    if (outputQueue != null) {
                        outputQueue.offer(mat.clone());
                    }
                }

                // 释放资源
                preprocessed.release();

            } catch (Exception e) {
                if (running) { // 仅在运行状态时打印错误
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 预处理图像
     * 缩放 -> 640 * 640
     * 颜色空间转换
     * 归一化
     */
    private Mat preprocessImage(Mat src) {
        Mat resized = new Mat();
        resize(src, resized, new Size(inputWidth, inputHeight));

        Mat rgb = new Mat();
        cvtColor(resized, rgb, COLOR_BGR2RGB);  // BGR转RGB

        rgb.convertTo(rgb, CV_32F, 1.0 / 255.0, 0);
        resized.release();
        return rgb;
    }

    /**
     * 将Mat HWC格式转换为CHW格式的数组
     */
    private float[] matToCHW(Mat mat) {
        int channels = mat.channels();
        int width = mat.cols();
        int height = mat.rows();

        float[] chw = new float[channels * width * height];

        // 获取图像数据的直接缓冲区
        FloatBuffer buffer = mat.createBuffer();

        // 按CHW格式重组数据
        for (int c = 0; c < channels; c++) {
            for (int y = 0; y < height; y++) {
                for (int x = 0; x < width; x++) {
                    // 计算在原始HWC格式中的索引
                    int hwcIndex = y * width * channels + x * channels + c;
                    // 计算在目标CHW格式中的索引
                    int chwIndex = c * width * height + y * width + x;
                    chw[chwIndex] = buffer.get(hwcIndex);
                }
            }
        }

        return chw;
    }

    private List<Detection> postprocessResult(OrtSession.Result result, int originalWidth, int originalHeight) throws OrtException {
        OnnxValue output = result.get(0);
        // 输出张量形状是 [1, 6, 8400]
        float[][][] detections = (float[][][]) ((OnnxTensor) output).getValue();

        List<Detection> detectionList = new ArrayList<>();


        // detections[batch][feature][box_index]
        // batch=1
        // feature=6 (cx, cy, w, h, confidence, classId)
        // box_index=8400

        int boxCount = detections[0][0].length;  // 8400 检测框数量

        for (int i = 0; i < boxCount; i++) {
            float confidence = detections[0][4][i];
            if (confidence < confidenceThreshold) continue;
            float cx = detections[0][0][i];
            float cy = detections[0][1][i];
            float w = detections[0][2][i];
            float h = detections[0][3][i];
            int classId = (int) detections[0][5][i];
//            System.out.println("检测结果：" + cx + ", " + cy + ", " + w + ", " + h + ", " + confidence + ", " + classId);
            // 转换为原图坐标
            float x1 = (cx - w / 2) * originalWidth / inputWidth;
            float y1 = (cy - h / 2) * originalHeight / inputHeight;
            float x2 = (cx + w / 2) * originalWidth / inputWidth;
            float y2 = (cy + h / 2) * originalHeight / inputHeight;
//            System.out.println("✅ 检测结果：" + x1 + ", " + y1 + ", " + x2 + ", " + y2 + ", " + confidence + ", " + classId);
            detectionList.add(new Detection(x1, y1, x2, y2, confidence, classId));
        }

        return applyNMS(detectionList, nmsThreshold);
    }


    private List<Detection> applyNMS(List<Detection> detections, float threshold) {
        detections.sort((a, b) -> Float.compare(b.confidence, a.confidence));
        List<Detection> result = new ArrayList<>();

        while (!detections.isEmpty()) {
            Detection best = detections.remove(0);
            result.add(best);

            Iterator<Detection> it = detections.iterator();
            while (it.hasNext()) {
                Detection detection = it.next();
                if (calculateIoU(best, detection) > threshold) {
                    it.remove();
                }
            }
        }

        return result;
    }

    private float calculateIoU(Detection a, Detection b) {
        float x1 = Math.max(a.x1, b.x1);
        float y1 = Math.max(a.y1, b.y1);
        float x2 = Math.min(a.x2, b.x2);
        float y2 = Math.min(a.y2, b.y2);

        float intersectionArea = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
        float unionArea = (a.x2 - a.x1) * (a.y2 - a.y1) +
                (b.x2 - b.x1) * (b.y2 - b.y1) -
                intersectionArea;

        return intersectionArea / unionArea;
    }

    // 内部类：检测结果
    public static class Detection {
        public float x1, y1, x2, y2;
        public float confidence;
        public int classId;

        public Detection(float x1, float y1, float x2, float y2, float confidence, int classId) {
            this.x1 = x1;
            this.y1 = y1;
            this.x2 = x2;
            this.y2 = y2;
            this.confidence = confidence;
            this.classId = classId;
        }
    }
    /**
     * 停止处理器线程
     */
    public void stop() {
        running = false;
        // 中断线程（可选，用于唤醒阻塞的take()）
        inputQueue.clear();
    }
}